/*
  Author: Chen Haodong (shinkunyachen@outlook.com)
  -----
  Copyright (c)2021 - present Chen Haodong. All Rights Reserved.
 */

#include "checkers/BranchChecker/BranchAnalyzer.h"
#include <clang/AST/Type.h>
#include <queue>
void BranchAnalyzer::analyze(ASTFile *file)
{
    auto ast_functions = file->getFunctions();
    for (auto f : ast_functions)
    {
        clang::FunctionDecl *func = manager->getFunctionDecl(f);
        //func->dump();
        LCFG lcfg(func);
        std::vector<int> costs(lcfg.basic_block_size());
        for (auto i = lcfg.basic_block_begin(), e = lcfg.basic_block_end();
             i != e; ++i)
        {
            LCFGBasicBlock *bb = *i;
            costs[bb->id()] = calc_basic_block_weight(func, bb);
        }
        calc_basic_block_costs(lcfg, costs);
        for (auto i = lcfg.element_begin(), e = lcfg.element_end(); i != e; ++i)
        {
            if ((*i)->kind() == LCFGElement::Statement)
            {
                const LCFGStatement *lcfg_stmt =
                    static_cast<const LCFGStatement *>(*i);
                const clang::Stmt *stmt = lcfg_stmt->stmt();
                switch (stmt->getStmtClass())
                {
                case clang::Stmt::IfStmtClass:
                {
                    double probability = calc_cond_probability(lcfg_stmt);
                    if (probability >= 0.0)
                        report(func, stmt, probability);
                    else
                    {
                        report(func, stmt, calc_cost_probability(lcfg_stmt, costs));
                    }
                }
                break;
                default:
                    break;
                }
            }
        }
    }
}

int BranchAnalyzer::calc_basic_block_weight(const clang::FunctionDecl *caller, const LCFGBasicBlock *bb)
{
    int cost = 0;
    for (auto i = bb->element_begin(), end = bb->element_end(); i != end; ++i)
    {
        int elm_cost = 0;
        LCFGElement *elm = (*i);
        if (elm->kind() == LCFGElement::Statement)
        {
            elm_cost = is_in_critical_path(caller, static_cast<LCFGStatement *>(elm)->stmt()) ? 20 : 1;
        }
        auto p_it = elm->plane_begin();
        ++p_it;
        for (auto endp = elm->plane_end(); p_it != endp; ++p_it)
        {
            elm_cost *= 10;
        }
        cost += elm_cost;
    }
    return cost;
}
void BranchAnalyzer::calc_basic_block_costs(const LCFG &lcfg,
                                            std::vector<int> &costs)
{
    std::vector<int> outdegrees(lcfg.basic_block_size(), 0);
    std::queue<LCFGBasicBlock *> queue;
    for (auto i = lcfg.basic_block_begin(), e = lcfg.basic_block_end(); i != e;
         ++i)
    {
        LCFGBasicBlock *bb = *i;
        for (auto j = bb->pred_begin(), je = bb->pred_end(); j != je; ++j)
        {
            ++outdegrees[(*j)->id()];
        }
    }
    for (auto i = lcfg.basic_block_begin(), e = lcfg.basic_block_end(); i != e;
         ++i)
    {
        if (outdegrees[(*i)->id()] == 0)
            queue.push(*i);
    }
    while (!queue.empty())
    {
        LCFGBasicBlock *bb = queue.front();
        queue.pop();
        for (auto i = bb->succ_begin(), e = bb->succ_end(); i != e; ++i)
        {
            costs[bb->id()] += costs[(*i)->id()];
        }
        for (auto i = bb->pred_begin(), e = bb->pred_end(); i != e; ++i)
        {
            --outdegrees[(*i)->id()];
            if (outdegrees[(*i)->id()] == 0)
            {
                queue.push(*i);
            }
        }
    }
}

double BranchAnalyzer::calc_cost_probability(const LCFGStatement *cond,
                                             const std::vector<int> &costs)
{
    const clang::Stmt *stmt = cond->stmt();
    const clang::Stmt *then_stmt = nullptr;
    switch (stmt->getStmtClass())
    {
    case clang::Stmt::IfStmtClass:
        then_stmt = llvm::cast<clang::IfStmt>(stmt)->getThen();
        break;
    default:
        break;
    }
    if (!then_stmt)
        return -1.0;
    if (then_stmt->getStmtClass() == clang::Stmt::CompoundStmtClass)
        then_stmt = llvm::cast<clang::CompoundStmt>(then_stmt)->body_front();
    if (!then_stmt)
        return -1.0;
    int cur_cost = 0;
    int total_cost = 0;
    const LCFGBasicBlock *cond_bb = cond->basic_block();
    for (auto i = cond_bb->succ_begin(), e = cond_bb->succ_end(); i != e; ++i)
    {
        const LCFGBasicBlock *succ = (*i);
        total_cost += costs[succ->id()];
        if (cur_cost == 0)
        {
            for (auto j = succ->element_begin(), je = succ->element_end(); j != je;
                 ++j)
            {
                if ((*j)->kind() == LCFGElement::Statement)
                    if ((static_cast<const LCFGStatement *>(*j))->stmt() == then_stmt)
                        cur_cost = costs[(*j)->basic_block()->id()];
            }
        }
    }
    return (double)cur_cost / (double)total_cost;
}
double BranchAnalyzer::calc_cond_probability(const LCFGStatement *cond)
{
    const clang::Stmt *stmt = cond->stmt();
    const clang::Expr *cond_expr = nullptr;
    switch (stmt->getStmtClass())
    {
    case clang::Stmt::IfStmtClass:
        cond_expr = llvm::cast<clang::IfStmt>(stmt)->getCond();
        break;
    default:
        break;
    }
    if (cond_expr)
    {
        return calc_cond_probability_impl(cond_expr);
    }
    return -1.0;
}
double BranchAnalyzer::calc_cond_probability_impl(const clang::Expr *expr)
{
    if (expr->getStmtClass() == clang::Stmt::BinaryOperatorClass)
    {
        auto b_expr = llvm::cast<clang::BinaryOperator>(expr);
        switch (b_expr->getOpcode())
        {
        case clang::BinaryOperatorKind::BO_LOr:
        {
            double lhs_prob = calc_cond_probability_impl(b_expr->getLHS());
            double rhs_prob = calc_cond_probability_impl(b_expr->getRHS());
            if (lhs_prob < 0 || rhs_prob < 0)
                return -1.0;
            else
                return 1.0 - (lhs_prob * rhs_prob);
        }
        case clang::BinaryOperatorKind::BO_LAnd:
        {
            double lhs_prob = calc_cond_probability_impl(b_expr->getLHS());
            double rhs_prob = calc_cond_probability_impl(b_expr->getRHS());
            if (lhs_prob < 0 || rhs_prob < 0)
                return -1.0;
            else
                return lhs_prob * rhs_prob;
        }
        case clang::BinaryOperatorKind::BO_LT:
        case clang::BinaryOperatorKind::BO_GT:
        case clang::BinaryOperatorKind::BO_LE:
        case clang::BinaryOperatorKind::BO_GE:
        case clang::BinaryOperatorKind::BO_EQ:
        case clang::BinaryOperatorKind::BO_NE:
            return _condition_table.probability(b_expr);
        }
    }
    else if (expr->getStmtClass() == clang::Stmt::UnaryOperatorClass)
    {
        auto u_expr = llvm::cast<clang::UnaryOperator>(expr);
        if (u_expr->getOpcode() == clang::UnaryOperatorKind::UO_LNot)
        {
            double prob = calc_cond_probability_impl(u_expr->getSubExpr());
            return prob < 0 ? prob : 1.0 - prob;
        }
    }
    else
    {
        return _condition_table.probability(expr);
    }
    return -1.0;
}
template <typename _Pred>
bool BranchAnalyzer::contain_expr(const clang::Expr *expr, _Pred &&predicate)
{
    if (predicate(expr))
    {
        return true;
    }
    if (expr->getStmtClass() >= clang::Stmt::firstCastExprConstant &&
        expr->getStmtClass() <= clang::Stmt::lastCastExprConstant)
    {
        return contain_expr(llvm::cast<clang::CastExpr>(expr)->getSubExpr(),
                            predicate);
    }
    else
    {
        switch (expr->getStmtClass())
        {
        case clang::Stmt::BinaryOperatorClass:
        {
            return contain_expr(llvm::cast<clang::BinaryOperator>(expr)->getLHS(),
                                predicate) ||
                   contain_expr(llvm::cast<clang::BinaryOperator>(expr)->getRHS(),
                                predicate);
            break;
        }
        case clang::Stmt::UnaryOperatorClass:
        {
            return contain_expr(llvm::cast<clang::UnaryOperator>(expr)->getSubExpr(),
                                predicate);
            break;
        }
        default:
            break;
        }
    }
    return false;
}
void BranchAnalyzer::report(clang::FunctionDecl *func_decl,
                            const clang::Stmt *stmt, double probability)
{
    _results.insert(std::make_pair(stmt, probability));
}
void BranchAnalyzer::load_critical_path()
{
    std::unordered_map<std::string, std::string> critical_path_cfg =
        configure->getOptionBlock("CriticalPaths");
    for (auto i = critical_path_cfg.begin(), e = critical_path_cfg.end(); i != e; ++i)
    {
        std::string &chain = i->second;
        size_t caller_b = 0;
        size_t caller_e = chain.find(',');
        size_t callee_b, callee_e;
        while (1)
        {
            if (caller_e == chain.npos || caller_e >= chain.length())
                break;
            callee_b = caller_e + 1;
            callee_e = chain.find(',', callee_b);
            if (callee_e == chain.npos)
            {
                _critical_paths.insert(std::make_pair(chain.substr(caller_b, caller_e - caller_b), chain.substr(callee_b)));
            }
            else
            {
                _critical_paths.insert(std::make_pair(chain.substr(caller_b, caller_e - caller_b), chain.substr(callee_b, callee_e - callee_b)));
            }
            caller_b = callee_b;
            caller_e = callee_e;
        }
    }
}
bool BranchAnalyzer::is_in_critical_path(const clang::FunctionDecl *caller, const clang::Expr *expr)
{
    if (!expr)
        return false;
    expr = ConditionProbabilityTable::remove_cast(expr);
    switch (expr->getStmtClass())
    {
    case clang::Stmt::BinaryOperatorClass:
        if (is_in_critical_path(caller, llvm::cast<clang::BinaryOperator>(expr)->getLHS()) || is_in_critical_path(caller, llvm::cast<clang::BinaryOperator>(expr)->getRHS()))
            return true;
        break;
    case clang::Stmt::UnaryOperatorClass:
        if (is_in_critical_path(caller, llvm::cast<clang::UnaryOperator>(expr)->getSubExpr()))
            return true;
        break;
    case clang::Stmt::CallExprClass:
    {
        const clang::FunctionDecl *callee = llvm::cast<clang::CallExpr>(expr)->getDirectCallee();
        if (callee)
        {
            std::string callee_name = callee->getQualifiedNameAsString();
            auto rge = _critical_paths.equal_range(caller->getQualifiedNameAsString());
            for (auto i = rge.first; i != rge.second; ++i)
            {
                if (i->second == callee_name)
                    return true;
            }
        }
        break;
    }
    default:
        break;
    }
    return false;
}
bool BranchAnalyzer::is_in_critical_path(const clang::FunctionDecl *caller, const clang::Stmt *stmt)
{
    if (!stmt)
        return false;
    switch (stmt->getStmtClass())
    {
    case clang::Stmt::CompoundStmtClass:
        for (auto it = llvm::cast<clang::CompoundStmt>(stmt)->body_begin(), end = llvm::cast<clang::CompoundStmt>(stmt)->body_end(); it != end; ++it)
        {
            if (is_in_critical_path(caller, *it))
                return true;
        }
        break;
    case clang::Stmt::DeclStmtClass:
    {
        const clang::DeclStmt *decl_stmt = llvm::cast<clang::DeclStmt>(stmt);
        if (decl_stmt->isSingleDecl())
        {
            const clang::Decl *decl = decl_stmt->getSingleDecl();
            if (decl && decl->getKind() == clang::Decl::Var)
            {
                if (is_in_critical_path(caller, llvm::cast<clang::VarDecl>(decl)->getInit()))
                    return true;
            }
        }
        break;
    }
    case clang::Stmt::DoStmtClass:
        if (is_in_critical_path(caller, llvm::cast<clang::DoStmt>(stmt)->getCond()))
            return true;
        break;
    case clang::Stmt::ForStmtClass:
        if (is_in_critical_path(caller, llvm::cast<clang::ForStmt>(stmt)->getInit()) || is_in_critical_path(caller, llvm::cast<clang::ForStmt>(stmt)->getCond()) || is_in_critical_path(caller, llvm::cast<clang::ForStmt>(stmt)->getInc()))
            return true;
        break;
    case clang::Stmt::WhileStmtClass:
        if (is_in_critical_path(caller, llvm::cast<clang::WhileStmt>(stmt)->getCond()))
            return true;
        break;
    case clang::Stmt::IfStmtClass:
        if (is_in_critical_path(caller, llvm::cast<clang::IfStmt>(stmt)->getCond()))
            return true;
        break;
    case clang::Stmt::ReturnStmtClass:
        if (is_in_critical_path(caller, llvm::cast<clang::ReturnStmt>(stmt)->getRetValue()))
            return true;
        break;
    case clang::Stmt::SwitchStmtClass:
        if (is_in_critical_path(caller, llvm::cast<clang::SwitchStmt>(stmt)->getBody()))
            return true;
        break;
    default:
        if (stmt->getStmtClass() >= clang::Stmt::firstExprConstant && stmt->getStmtClass() <= clang::Stmt::lastExprConstant)
            return is_in_critical_path(caller, llvm::cast<clang::Expr>(stmt));
        break;
    }
    return false;
}

const clang::Expr *
ConditionProbabilityTable::remove_cast(const clang::Expr *expr)
{
    if (expr->getStmtClass() >= clang::Stmt::firstCastExprConstant &&
        expr->getStmtClass() <= clang::Stmt::lastCastExprConstant)
    {
        return remove_cast(llvm::cast<clang::CastExpr>(expr)->getSubExpr());
    }
    return expr;
}
ConditionProbabilityTable::ConditionProbabilityTable() { init_guessers(); }
double ConditionProbabilityTable::probability(const clang::Expr *expr)
{
    Condition cond;
    if (!extract_cond(expr, cond))
        return -1.0;
    return guess(cond);
}
bool ConditionProbabilityTable::extract_cond(const clang::Expr *expr,
                                             Condition &cond)
{
    if (expr->getStmtClass() == clang::Stmt::BinaryOperatorClass)
    {
        auto b_expr = llvm::cast<clang::BinaryOperator>(expr);
        switch (b_expr->getOpcode())
        {
        case clang::BinaryOperatorKind::BO_LT:
        case clang::BinaryOperatorKind::BO_GT:
        case clang::BinaryOperatorKind::BO_LE:
        case clang::BinaryOperatorKind::BO_GE:
        case clang::BinaryOperatorKind::BO_EQ:
        case clang::BinaryOperatorKind::BO_NE:
        {
            const clang::Expr *lhs = remove_cast(b_expr->getLHS());
            const clang::Expr *rhs = remove_cast(b_expr->getRHS());
            if (lhs->getStmtClass() == clang::Stmt::GNUNullExprClass ||
                lhs->getStmtClass() == clang::Stmt::CXXNullPtrLiteralExprClass ||
                lhs->getStmtClass() == clang::Stmt::IntegerLiteralClass)
            {
                cond.lhs = rhs;
                cond.rhs = lhs;
                switch (b_expr->getOpcode())
                {
                case clang::BinaryOperatorKind::BO_LT:
                    cond.op = ComparisonOp::GT;
                    break;
                case clang::BinaryOperatorKind::BO_GT:
                    cond.op = ComparisonOp::LT;
                    break;
                case clang::BinaryOperatorKind::BO_LE:
                    cond.op = ComparisonOp::GE;
                    break;
                case clang::BinaryOperatorKind::BO_GE:
                    cond.op = ComparisonOp::LE;
                    break;
                case clang::BinaryOperatorKind::BO_EQ:
                    cond.op = ComparisonOp::EQ;
                    break;
                case clang::BinaryOperatorKind::BO_NE:
                    cond.op = ComparisonOp::NEQ;
                    break;
                }
            }
            else
            {
                cond.lhs = lhs;
                cond.rhs = rhs;
                switch (b_expr->getOpcode())
                {
                case clang::BinaryOperatorKind::BO_LT:
                    cond.op = ComparisonOp::LT;
                    break;
                case clang::BinaryOperatorKind::BO_GT:
                    cond.op = ComparisonOp::GT;
                    break;
                case clang::BinaryOperatorKind::BO_LE:
                    cond.op = ComparisonOp::LE;
                    break;
                case clang::BinaryOperatorKind::BO_GE:
                    cond.op = ComparisonOp::GE;
                    break;
                case clang::BinaryOperatorKind::BO_EQ:
                    cond.op = ComparisonOp::EQ;
                    break;
                case clang::BinaryOperatorKind::BO_NE:
                    cond.op = ComparisonOp::NEQ;
                    break;
                }
            }
            return true;
        }
        default:
            return false;
        }
    }
    else
    {
        cond.lhs = remove_cast(expr);
        cond.op = ComparisonOp::NEQ;
        cond.rhs = nullptr;
        return true;
    }
    return false;
}
double ConditionProbabilityTable::guess(const Condition &cond)
{
    for (auto g : _guessers)
    {
        double result = g(cond);
        if (result >= 0.0)
            return result;
    }
    return -1.0;
}
void ConditionProbabilityTable::init_guessers()
{
    _guessers.push_back([](const Condition &cond) -> double {
        const clang::Type *type = nullptr;
        if (cond.lhs->getStmtClass() == clang::Stmt::DeclRefExprClass)
        {
            type = llvm::cast<clang::DeclRefExpr>(cond.lhs)
                       ->getDecl()
                       ->getType()
                       .getTypePtr();
        }
        if (cond.lhs->getStmtClass() == clang::Stmt::CallExprClass && llvm::cast<clang::CallExpr>(cond.lhs)->getDirectCallee())
        {
            type = llvm::cast<clang::CallExpr>(cond.lhs)->getDirectCallee()->getReturnType().getTypePtr();
        }
        if (type && type->isPointerType())
        {
            if (cond.op == ComparisonOp::EQ)
                return 0.1;
            else if (cond.op == ComparisonOp::NEQ)
                return 0.9;
        }
        return -1.0;
    });
    _guessers.push_back([](const Condition &cond) -> double {
        if (cond.rhs &&
            cond.rhs->getStmtClass() == clang::Stmt::IntegerLiteralClass &&
            llvm::cast<clang::IntegerLiteral>(cond.rhs)->getValue().isNullValue())
        {
            switch (cond.op)
            {
            case ComparisonOp::GE:
            case ComparisonOp::GT:
                return 0.9;
            case ComparisonOp::LT:
            case ComparisonOp::LE:
                return 0.1;
                case ComparisonOp::EQ:
                    return 0.9;
                    case ComparisonOp::NEQ:
                    return 0.1;
            }
        }
        return -1.0;
    });
    _guessers.push_back([](const Condition &cond) -> double {
        if (cond.rhs && cond.rhs->getStmtClass() == clang::Stmt::DeclRefExprClass)
        {
            std::string identifier =
                llvm::cast<clang::DeclRefExpr>(cond.rhs)->getDecl()->getName().str();
            if (identifier.find("sz") != identifier.npos ||
                identifier.find("size") != identifier.npos ||
                identifier.find("cnt") != identifier.npos ||
                identifier.find("count") != identifier.npos ||
                identifier.find("len") != identifier.npos ||
                identifier.find("length") != identifier.npos)
                switch (cond.op)
                {
                case ComparisonOp::GE:
                case ComparisonOp::GT:
                    return 0.1;
                case ComparisonOp::LT:
                case ComparisonOp::LE:
                    return 0.9;
                }
            if(identifier.find("OK") != identifier.npos||
            identifier.find("Ok") != identifier.npos||
            identifier.find("ok") != identifier.npos){
                switch (cond.op)
                {
                case ComparisonOp::EQ:
                    return 0.9;
                case ComparisonOp::NEQ:
                    return 0.1;
                }
            }
            if(identifier.find("Fail") != identifier.npos||
            identifier.find("fail") != identifier.npos||
            identifier.find("FAIL") != identifier.npos){
                switch (cond.op)
                {
                case ComparisonOp::EQ:
                    return 0.1;
                case ComparisonOp::NEQ:
                    return 0.9;
                }
            }
            if(identifier.find("ERR") != identifier.npos||
            identifier.find("Err") != identifier.npos||
            identifier.find("err") != identifier.npos){
                if(identifier.find("No") != identifier.npos||
                identifier.find("NO") != identifier.npos||
                identifier.find("no") != identifier.npos){
                    switch (cond.op)
                    {
                    case ComparisonOp::EQ:
                        return 0.9;
                    case ComparisonOp::NEQ:
                        return 0.1;
                    }
                }
                else{
                    switch (cond.op)
                    {
                    case ComparisonOp::EQ:
                        return 0.1;
                    case ComparisonOp::NEQ:
                        return 0.9;
                    }
                }
            }
        }
        return -1.0;
    });
    _guessers.push_back([](const Condition &cond) -> double {
        const clang::CallExpr *call_expr = nullptr;
        if (cond.lhs->getStmtClass() == clang::Stmt::CallExprClass)
        {
            call_expr = llvm::cast<clang::CallExpr>(cond.lhs);
        }
        if (cond.rhs && cond.rhs->getStmtClass() == clang::Stmt::CallExprClass)
        {
            call_expr = llvm::cast<clang::CallExpr>(cond.rhs);
        }
        if (!call_expr)
            return -1.0;
        const clang::FunctionDecl *func_decl = call_expr->getDirectCallee();
        if (!func_decl)
            return -1.0;
        std::string function_name = func_decl->getQualifiedNameAsString();
        if(function_name == "__builtin_expect"){
            if(func_decl->param_size() == 2){
                const clang::IntegerLiteral *literal = llvm::cast<clang::IntegerLiteral>(ConditionProbabilityTable::remove_cast(call_expr->getArg(1)));
                if(literal->getValue().isNullValue())
                    return 0.0;
                else
                    return 1.0;
            }
        }
        if (function_name == "feof")
            switch (cond.op)
            {
            case ComparisonOp::EQ:
                return 0.9;
            case ComparisonOp::NEQ:
                return 0.1;
            }

        return -1.0;
    });
}